
package ploy.compiler;

import java.util.*;

import org.objectweb.asm.*;
import org.objectweb.asm.util.*;
import org.objectweb.asm.commons.*;

import ploy.lang.*;

/**
 * Lambda represents AST nodes of the form (lambda (...) ...).  As functions
 * are the basic unit of computability in Ploy, compiling a function
 * is the main entrypoint into compiling expressions, which is where most
 * compilation occurs.
 */
class Lambda extends AST implements Comparable<Lambda> {
		
    /* Useful type names. */
    public static final String[] invokeExceptions = 
				{ "ploy/lang/RuntimeError" };

    /* The number of arguments before they get passed in a list. */
    public static final int maxInlineArgs = 4;
		
    /* List of formal arguments to the lambda. */
    private List<Symbol> args;
		
    /* The body of the function. */
    private AST exp;

    /* Lambdas memoize their variable location context. */
    VarLocationTable varLocContext;
		
    /* The name for the generated class representing instances
     * of this function. */
    public Symbol generatedName;
		
    public Lambda( List<Symbol> a, AST e ) {
	args = a;
	exp = e;
	generatedName = Symbol.genSym( );
    }
    
    public String toString( ) {
	String toReturn = "(lambda (";
	for ( Iterator<Symbol> i = args.iterator(); i.hasNext(); ) {
	    toReturn += i.next().toString();
	    if ( i.hasNext() ) 
		toReturn += " ";
	}
	toReturn += ") " + exp.toString() + ")";
	return toReturn;
    }

    /**
     * To calculate free variables of a lambda, calculate the free varaibles
     * of the body, and then remove the argument.
     */
    protected TreeSet<Symbol> calculateFreeVars( ) {
	TreeSet<Symbol> s = (TreeSet<Symbol>)exp.freeVars().clone();
	for ( Iterator<Symbol> i = args.iterator(); i.hasNext(); ) {
	    s.remove( i.next() );
	}
	return s;
    }
    
    /* An advantage of having a generated name is that we can use it for
     * quick string operations. */
    public int hashCode( ) {
	return generatedName.hashCode( );
    }
    
    public int compareTo( Lambda o ) {
	return generatedName.compareTo( o.generatedName );
    }
		
    public boolean equals( Object o ) {
	return ( ( o instanceof Lambda) && 
		 ( generatedName == ((Lambda)o).generatedName ) );
    }
    
    /**
     * The lambdas of this AST is itself and all of the lambdas in the body.
     */
    public TreeSet<Lambda> allLambdas( ) {
	TreeSet<Lambda> s = 
	    (TreeSet<Lambda>)exp.allLambdas().clone(); /* necessary? */
	s.add( this );
	return s;
    }
    
    /**
     * The JVM type of the JVM function used to generate instances of this
     * Ploy function.
     */
    public String constructorType;
		
    /**
     * The JVM name and type of the JVM function used to invoke this function.
     */
    public String invokeType;
    public String invokeName;

    public static String invokeType( int n ) {
	String invokeType;
	if ( n > maxInlineArgs ) {
	    invokeType = "(" + DynamicObject.jvmType + ")" + 
		DynamicObject.jvmType;
	} else {						
	    invokeType = "(";
	    for ( int j = 0; j < n; j++ ) {
		invokeType += DynamicObject.jvmType;
	    }
	    invokeType += ")" + DynamicObject.jvmType;
	}
	return invokeType;
    }
	
    public static String invokeName( int n ) {
	String invokeName;
	if ( n > maxInlineArgs ) {
	    invokeName = "invokeN";
	} else {						
	    invokeName = "invoke" + Integer.toString( n );
	}
	return invokeName;
    }

    /**
     * Generate the bytecode representing the compiled JVM class representing
     * instances of this JVM function.
     */
    public void generateClass( ClassVisitor cv ) {

	/* Generate the class header. */
        cv.visit( Opcodes.V1_5, Opcodes.ACC_PUBLIC, 
		  generatedName.toString(), null,
                  Closure.jvmClass, null );
	
	/* For each free variable that requires global storage,
	 * Generate a field in the closure. */
        Iterator<Symbol> i = freeVars().iterator();
        while( i.hasNext() ) {
            Symbol s = i.next();
	    VarLocation v = varLocContext.get( s );
            if ( v.needsLocalStorage() ) {
		cv.visitField(Opcodes.ACC_PRIVATE, s.toString(), 
			      DynamicObject.jvmType,
			      null, null).visitEnd();
	    }
        }
	
	/* For each field in the closure, add it to the constructor
	 * arugments. */
        constructorType = "(";
        i = freeVars().iterator();
        while( i.hasNext() ) {
            Symbol s = i.next();
	    VarLocation v = varLocContext.get( s );
            if ( v.needsLocalStorage() ) {
		constructorType += DynamicObject.jvmType;
	    }
        }
        constructorType += ")V";
	
	/* Generate the constructor. */
        MethodVisitor constructorMV =
            cv.visitMethod(Opcodes.ACC_PUBLIC, "<init>", constructorType, 
                           null, null);
        constructorMV.visitCode();
	constructorMV.visitVarInsn(Opcodes.ALOAD, 0);
	constructorMV.visitIntInsn(Opcodes.SIPUSH, args.size());
	constructorMV.visitMethodInsn(Opcodes.INVOKESPECIAL, Closure.jvmClass,
				      "<init>", "(I)V");
	
	/* For each argument to the contructor, store it in its appropriate
	 * field. */
	i = freeVars().iterator();
        int n = 0;
        int consMaxS = 1;
        while( i.hasNext() ) {
            Symbol s = i.next();
	    VarLocation v = varLocContext.get( s );
	    if ( v.needsLocalStorage() ) {
		n++;
		consMaxS++;
		constructorMV.visitVarInsn(Opcodes.ALOAD, 0);
		constructorMV.visitVarInsn(Opcodes.ALOAD, n);
		constructorMV.visitFieldInsn(Opcodes.PUTFIELD, 
					     generatedName.toString(),
					     s.toString(), 
					     DynamicObject.jvmType);
	    }
        }
        constructorMV.visitInsn(Opcodes.RETURN);
        constructorMV.visitMaxs( 2, consMaxS );
        constructorMV.visitEnd();
	
	/* The invoke method has a signature that can only be computed
	 * once the number of arguments are know. */	
	int numOfArgs = args.size();
	invokeType = invokeType( numOfArgs );
	invokeName = invokeName( numOfArgs );

	/* We also need an "invoke" method.  Compiling this delegates to its
	 * body. */
        MethodVisitor invokeMV = 
            cv.visitMethod(Opcodes.ACC_PUBLIC, invokeName, 
			   invokeType, null, invokeExceptions);
        FrameSize invokeFS = exp.generateBytecode( invokeMV );
        invokeMV.visitInsn(Opcodes.ARETURN);
        invokeMV.visitMaxs( invokeFS.maxS, 1 + numOfArgs );
        invokeMV.visitEnd();
        
        cv.visitEnd();
    }
    
    /**
     * A compiled lambda expression instantiates its associated class.
     */
    public FrameSize generateBytecode( MethodVisitor mv ) {
        FrameSize curS = new FrameSize( 2, 0 );
        FrameSize maxS = curS;
        mv.visitTypeInsn(Opcodes.NEW, generatedName.toString());
        mv.visitInsn(Opcodes.DUP);
        Iterator<Symbol> i = freeVars().iterator();
        while( i.hasNext() ) {
            Symbol s = i.next();
	    VarLocation v = varLocContext.get(s);
	    if ( v.needsLocalStorage() ) {
		FrameSize subS = v.generateBytecode(mv);
		maxS = FrameSize.max( maxS, curS.add( subS ) );
		curS = curS.addS( 1 );
	    }
        }
        mv.visitMethodInsn(Opcodes.INVOKESPECIAL, generatedName.toString(),
                           "<init>", constructorType);
        return maxS;
    }    
		
    /**
     * Locating variables in the lambda does four things:
     *
     * 1. memoizes the table
     * 2. puts the free variables on the table as "current-closure" vars
     * 3. puts the argument variable on the table
     * 4. locates the variables in the body.
     */
    public void locateVariables( VarLocationTable t ) {
	varLocContext = t;
        VarLocationTable m = t.clone( );
	Iterator<Symbol> freeVI = freeVars().iterator();
        while( freeVI.hasNext() ) {
            Symbol s = freeVI.next();
	    VarLocation v = varLocContext.get( s );
            if ( v.needsLocalStorage() ) {
		m.put( s, new CurrentClosureVarLocation
		       ( generatedName.toString(), s.toString() ) );
	    }
        }
	Iterator<Symbol> argI = args.iterator();
	int argPos = 1;
	while ( argI.hasNext() ) {
	    m.put( argI.next(), new ArgumentLocation( argPos++ ) );
	}
	exp.locateVariables( m );
	return;
    }
    
    
}